import dashscope
from http import HTTPStatus
import json
import sys
# 设置控制台输出编码为 UTF-8
sys.stdout.reconfigure(encoding='utf-8')

def check_status(response, operation):
    if response.status_code == HTTPStatus.OK:
        print(f"{operation} 成功。")
        return True
    else:
        print(f"{operation} 失败。状态码：{response.status_code}，错误码：{response.code}，错误信息：{response.message}")
        sys.exit(response.status_code)

# 1. 创建绘画助手
def create_painting_assistant():
    return dashscope.Assistants.create(
        model='qwen-max',  # 模型列表：https://help.aliyun.com/zh/model-studio/getting-started/models
        name='Art Maestro',
        description='AI助手，用于绘画和艺术知识',
        instructions='''提供绘画技巧、艺术史和创意指导的信息。
        使用工具进行研究和生成图像。''',
        tools=[
            {'type': 'quark_search', 'description': '用于研究艺术主题'},
            {'type': 'text_to_image', 'description': '用于创建视觉示例'}
        ]
    )

if __name__ == '__main__':
    # 创建绘画助手
    painting_assistant = create_painting_assistant()
    print(painting_assistant)
    check_status(painting_assistant, "助手创建")

    # 创建一个带有初始消息的新线程
    thread = dashscope.Threads.create(
        messages=[{
            'role': 'user',
            'content': '请帮我画一幅布偶猫的画。'
        }]
    )
    print(thread)
    check_status(thread, "线程创建")

    # 创建流式输出的运行
    run_iterator = dashscope.Runs.create(
        thread.id,
        assistant_id=painting_assistant.id,
        stream=True
    )

    # 迭代事件和消息
    print("处理请求中...")
    for event, msg in run_iterator:
        print(event)
        print(msg)

    # 检索并显示助手的响应
    messages = dashscope.Messages.list(thread.id)
    check_status(messages, "消息检索")

    print("\n助手的回应：")
    print(json.dumps(messages, ensure_ascii=False, default=lambda o: o.__dict__, sort_keys=True, indent=4))

# 提示: 这个脚本创建了一个流式输出的绘画助手，开始一段关于绘制布偶猫的对话，
# 并实时显示助手的回应。